#可以单独提前运算
#也可以被调用运算
#准备内容，波动率，放缩因子，系数计算；防追高数据等
from jili.core.printlog import print
import datetime
from gm import api
import pandas as pd
import baostock as bs
import sys
from jili.core.convert import str2datetime
from jili.gmlib import get_tradeday,get_tradeday_bydate
from jili.data.db import tradecal,get_calender
def str2float(x):
    if x.strip()=="":
        return 0.0
    else:
        return float(x)
def str2int(x):
    if x.strip()=="":
        return 0
    else:
        return int(x)
def date2datetime(x):
    return datetime.datetime.strptime(x,"%Y-%m-%d")
def time2datetime(x):
    #print(x)
    return datetime.datetime.strptime(x,"%Y%m%d%H%M%S%f")
def deal_df_str2float(k,field):
    for i in field:
        if i in k.columns:
            k[i]=k[i].apply(str2float)
def deal_df_str2int(k,field):
    for i in field:
        if i in k.columns:
            k[i]=k[i].apply(str2int)
def getk_bybaostock(bs,obj, field, start_date='1990-12-19', end_date=None, adjustflag="3"):
    now = datetime.datetime.now()
    if end_date is None:
        end_date = now.strftime("%Y-%m-%d")
    rst = []
    if isinstance(start_date,datetime.datetime):
        start_date=start_date.strftime("%Y-%m-%d")
    if isinstance(end_date,datetime.datetime):
        end_date=end_date.strftime("%Y-%m-%d")
    if len(start_date) == 8:
        start_date = start_date[:4] + "-" + start_date[4:6] + "-" + start_date[6:]
    if len(end_date) == 8:
        end_date = end_date[:4] + "-" + end_date[4:6] + "-" + end_date[6:]
    print("deal",obj,start_date,end_date,"d",adjustflag)
    rs = bs.query_history_k_data_plus(obj,
                                      field,
                                      start_date=start_date, end_date=end_date,
                                      frequency="d", adjustflag=adjustflag)
    if rs.error_code != "0":
        print(rs.error_code, rs.error_msg)
    else:
        #### 打印结果集 ####
        data_list = []
        while (rs.error_code == '0') & rs.next():
            # 获取一条记录，将记录合并在一起
            data_list.append(rs.get_row_data())
        kd = pd.DataFrame(data_list, columns=rs.fields)
        kd["obj"] = obj[-6:]
        kd["uptimekey"]=now
        deal_df_str2float(kd, ["open", "high", "low", "close", "preclose", "pctChg", "turn", "amount"])
        deal_df_str2int(kd, ["volume", "adjustflag", "tradestatus", "isST"])
        # kd["timekey"]=kd["date"].apply(date2datetime)
        k = kd.to_dict("records")
        for i in k:
            if "time" in i.keys():
                i["timekey"] = time2datetime(i["time"])
            else:
                i["timekey"] = date2datetime(i["date"])
            rst.append(i)
    print("get kline", "d", obj, len(rst))
    return rst
def get_bars_baosotock(objs,start,end):
    rst={}
    if objs:
        lg = bs.login()
        if lg.error_code!="0":
            print("baostock 获取数据失败",lg.error_code,lg.error_msg)
            sys.exit()
    field="date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST"

    for obj in objs:
        rst[obj]=getk_bybaostock(bs,obj, field, start_date=start, end_date=end, adjustflag="3")
    if objs:
        bs.logout()
    return rst
def xttime2datetime(i):
    return datetime.datetime.fromtimestamp(i/1000)
class predata:
    def __init__(self,objs=[],td=None,count=1):
        self.bars={}
        self.bars_byday = {}
        self.prebars = {}
        if td is None:
            td=datetime.datetime.now()
        self.ts=get_tradeday_bydate(td.strftime("%Y%m%d"),count=count)
        self.start=self.ts[0]
        self.end=self.ts[-1]
        self.count=count
        if objs:
            for obj in objs:
                bars0=self.deal_data2(obj,end=self.end,count=count)
                source="gm"
                if not bars0:
                    bars0=get_bars_baosotock([obj],self.start,self.end)
                    source = "baostock"
                self.bars[obj]=bars0
            # bars = xtdata.get_market_data(field_list=[], stock_list=objs, period='1d', start_time='',
            #                             end_time='', count=60, dividend_type='front_ratio', fill_data=True)
            # ds1 = xtdata.get_market_data(field_list=[], stock_list=objs, period='1d', start_time='',
            #                             end_time='', count=1, dividend_type='none', fill_data=True)
            for obj,v in self.bars.items():
                self.prebars[obj]=v[-1]
            print("数据加载完成，数据源：",source)
    def get_bars(self,obj):
        if obj not in self.bars.keys():
            bars0 = self.deal_data2(obj,end=self.end,count=self.count)
            # if not bars0:
            #     bars0 = get_bars_baosotock([obj], self.start, self.end)
            self.bars[obj]=bars0
            self.prebars[obj] = bars0[-1]
        return self.bars[obj]
    def get_bars_bydate(self,obj,start,end,dividend_type='front_ratio'):
        bars0 = self.deal_data([obj], start, end,dividend_type)
        return bars0[obj]
    def get_all_bars_bydate(self,start,end,dividend_type='front_ratio'):
        objs = xtdata.get_stock_list_in_sector(sector_name="沪深A股")
        bars0 = self.deal_data(objs, start, end,dividend_type)
        return bars0
    def load_all_k1d(self,product="stock"):
        if product == "index":
            objs = xtdata.get_stock_list_in_sector(sector_name='沪深指数')
        elif product == "fund":
            objs = xtdata.get_stock_list_in_sector(sector_name='沪深ETF')
        elif product == "bond":
            objs = xtdata.get_stock_list_in_sector(sector_name='沪深转债')
        else:
            objs = xtdata.get_stock_list_in_sector(sector_name="沪深A股")
        now=datetime.datetime.now()
        if now.hour<15:
            end=tradecal.getpretradeday(now)
        else:
            end=now
        xtdata.download_history_data2(objs, "1d", end_time=end.strftime("%Y%m%d"),callback=self.on_progress)
        print("增量更新完成",product,len(objs), "k1d")
    def save_dataset(self,baseurl,fqtype="bfq",product="stock"):
        from jili.dataset import dataset
        from copy import deepcopy
        if product == "index":
            objs = xtdata.get_stock_list_in_sector(sector_name='沪深指数')
        elif product == "fund":
            objs = xtdata.get_stock_list_in_sector(sector_name='沪深ETF')
        elif product == "bond":
            objs = xtdata.get_stock_list_in_sector(sector_name='沪深转债')
        else:
            objs = xtdata.get_stock_list_in_sector(sector_name="沪深A股")
        fqtype2xt = {"bfq": "none", "hfq": "back_ratio", "qfq": "front_ratio"}
        datatool = dataset(product=product, datatype="k1d", baseurl=baseurl, region="C", fqtype=fqtype,source="xt")
        start=datatool.get_maxdate("all")
        data0= {}
        if start:
            for i in datatool.get_data_bydate("all",start):
                obj=i["obj"]
                data0[obj]=i
            start0=tradecal.getnexttradeday(start)
            start0=start0.strftime("%Y%m%d")
        else:
            start0="20100101"
        now=datetime.datetime.now()
        if now.hour<15:
            end=tradecal.getpretradeday(now)
        else:
            end=now
        end=end.strftime("%Y%m%d")
        end0=str2datetime(start0)
        print("save",len(objs),start0,end)
        bars = xtdata.get_local_data(field_list=[], stock_list=objs, period="1d",
                                     start_time=start0, end_time=end,dividend_type=fqtype2xt[fqtype],
                                     fill_data=True)
        bars0 = {}
        isone = True
        for k, v in bars.items():
            if ".SZ" in k or ".SH" in k:
                isone = False
                obj=k.split(".")[0]
                v.dropna(inplace=True)
                for d, vv in v.to_dict("index").items():
                    d=str2datetime(d)
                    if d>end0:
                        end0=d
                    if obj not in bars0.keys():
                        bars0[obj] = {}
                    if d not in bars0[obj].keys():
                        bars0[obj][d] = {}
                    vv["timekey"] = d
                    vv["obj"]=obj
                    bars0[obj][d] = vv
            else:
                for obj, vv in v.to_dict("index").items():
                    obj=obj.split(".")[0]
                    for d, vvv in vv.items():
                        d = str2datetime(d)
                        if d > end0:
                            end0 = d
                        if obj not in bars0.keys():
                            bars0[obj] = {}
                        if d not in bars0[obj].keys():
                            bars0[obj][d] = {"obj":obj}
                        bars0[obj][d][k] = vvv
                        bars0[obj][d]["timekey"]=d
        rst = {}
        start0=str2datetime(start0)
        # objs1=[]
        # for i in objs:
        #     objs1.append(i.split(".")[0])
        if data0:
            noobjs=[]
            for obj0 in objs:
                if obj0 not in bars.keys():
                    obj = obj0.split(".")[0]
                    if obj in data0.keys():
                        noobjs.append(obj)
            for obj in noobjs:
                rst[obj]={}
                for d in get_calender(start0,end0):
                    b=deepcopy(data0[obj])
                    b["timekey"]=d
                    b['volume']=0
                    b["amount"]=0
                    rst[obj][d]=b
            for obj,v in bars0.items():
                l=list(v.keys())
                l.sort()
                if start0<l[0]:
                    if obj in data0.keys():
                        end01=tradecal.getpretradeday(l[0])
                        for d in get_calender(start0,end01):
                            b = deepcopy(data0[obj])
                            b["timekey"] = d
                            b['volume'] = 0
                            b["amount"] = 0
                            rst[obj][d] = b
                if obj not in rst.keys():
                    rst[obj]={}
                rst[obj].update(v)
            for obj,v in data0.items():
                if obj not in rst.keys():
                    rst[obj] = {}
                    for d in get_calender(start0, end0):
                        b = deepcopy(data0[obj])
                        b["timekey"] = d
                        b['volume'] = 0
                        b["amount"] = 0
                        rst[obj][d] = b
        else:
            rst=bars0


        alldata={}
        for obj, v in rst.items():
            datatool.save(obj.split(".")[0], v)
            print("save", obj, len(v))
            for d,vv in v.items():
                key=d.strftime("%Y%m")
                if key not in alldata.keys():
                    alldata[key]={}
                if d not in alldata[key].keys():
                    alldata[key][d]=[]
                alldata[key][d].append(vv)
        datatool.save_byall(alldata)
        print("save all")

    def covert_data_obj2date(self,data):
        rst={}
        for obj,v in data.items():
            if isinstance(v,list):
                for i in v:
                    d=i["timekey"]
                    if d not in rst.keys():
                        rst[d]={}
                    rst[d][obj]=i
            else:
                for d,i in v:
                    if d not in rst.keys():
                        rst[d]={}
                    rst[d][obj]=i
        return rst
    def on_progress(data):
        print(data)
    def on_progress0(data):
        print(data)
    def load_all_financial(self,start=None):
        objs = xtdata.get_stock_list_in_sector(sector_name="沪深A股")
        xtdata.download_financial_data2(objs, table_list=[], start_time='', end_time='', callback=self.on_progress0)
    def get_all_financial(self,end=None,key="Top10flowholder"):
        objs = xtdata.get_stock_list_in_sector(sector_name="沪深A股")
        if end:
            end=str2datetime(end).strftime("%Y%m%d")
        else:
            end=""
        bars=xtdata.get_financial_data(objs, table_list=[key], start_time='', end_time=end, report_type='announce_time')
        bars0={}
        for k, v in bars.items():
            for obj, vv in v.to_dict("index").items():
                for d, vvv in vv.items():
                    if obj not in bars0.keys():
                        bars0[obj] = {}
                    if d not in bars0[obj].keys():
                        bars0[obj][d] = {}
                    bars0[obj][d][k] = vvv
        rst={}
        for obj,v in bars0.items():
            for d,vv in v.items():
                vv["timekey"]=str2datetime(d)
            rst[obj]=list(v.values())
        return rst
    def deal_data(self,objs,start,end,dividend_type="front_ratio"):
        if start not in ["",None]:
            start=str2datetime(start)
        if end not in ["",None]:
            end=str2datetime(end)
        bars=xtdata.get_market_data(field_list=[], stock_list=objs, period='1d', start_time=start.strftime("%Y%m%d"),end_time=end.strftime("%Y%m%d"),dividend_type=dividend_type, fill_data=True)
        bars0 = {}
        for k, v in bars.items():
            for obj, vv in v.to_dict("index").items():
                for d, vvv in vv.items():
                    if obj not in bars0.keys():
                        bars0[obj] = {}
                    if d not in bars0[obj].keys():
                        bars0[obj][d] = {}
                    bars0[obj][d][k] = vvv
        rst={}
        for obj,v in bars0.items():
            for d,vv in v.items():
                vv["timekey"]=str2datetime(d)
            rst[obj]=list(v.values())
        return rst
    def deal_data2(self,obj,end,count=1,dividend_type=1,start=None):
        if end not in ["",None]:
            end=str2datetime(end)
        else:
            end=self.ts[-1]
        if start:
            start=str2datetime(start)
            bars = api.history(obj, '1d', start_time=start.strftime("%Y-%m-%d"), end_time=end.strftime("%Y-%m-%d"), adjust=dividend_type)
        else:
            bars=api.history_n(obj,'1d',count=count,end_time=end.strftime("%Y-%m-%d"),adjust=dividend_type)
        t=[]
        for i in bars:
            i["obj"]=obj
            i["timekey"]=i["bob"]
            t.append(i)
        return t
xtpredata=predata()
if __name__=="__main__":
    pass